//
// itertool - a polynormial and type-safe iterator framework in C++.
//
// Copyright (c) 2021-2025 Feng Cong
//
// Distributed under the MIT Software License
// See accompanying file LICENSE.txt or copy at
// https://opensource.org/licenses/MIT

#pragma once

#include <itertool/itertool_fwd.h>

#include <algorithm>
#include <iterator>
#include <memory>
#include <tuple>
#include <vector>

#if !defined(__cpp_fold_expressions)
#error "Zip requires fold expressions to work."
#endif

#define ITERTOOL_ZIP_REQUIRE_MESSAGE \
  "Zip/ZipLongest requires at least 2 input iterators"

namespace itertool {

template <typename... T>
using ZipTupleItem = std::tuple<T...>;

namespace detail {

template <typename TypeSeq, typename IndexSeq>
class ZipTupleIterable;

template <typename TypeSeq, typename IndexSeq>
class ZipLongestTupleIterable;

template <typename... T, size_t... Index>
class ZipLongestTupleIterable<TypeSequence<T...>, std::index_sequence<Index...>>
    : public Iterable<ZipTupleItem<T...>> {
 public:
  using ItemType = ZipTupleItem<T...>;

  ZipLongestTupleIterable(IteratorTuple<T...> data,
                          std::tuple<T...> default_value)
      : data_(data), default_value_(default_value) {}

  bool HasNext() override {  //
    return ITERTOOL_DETAIL_FoldExpr_AnyOf(std::get<Index>(data_).HasNext());
  }

  ItemType GetNow() override {  //
    return ItemType(MaybeGetNow<T, Index>()...);
  }

  void Advance() override {  //
    ITERTOOL_DETAIL_FoldExpr_ForEach(std::get<Index>(data_).MaybeAdvance());
  }

  bool IsFinite() override {  //
    return ITERTOOL_DETAIL_FoldExpr_AllOf(std::get<Index>(data_).IsFinite());
  }

  bool IsResetable() override {  //
    return ITERTOOL_DETAIL_FoldExpr_AllOf(std::get<Index>(data_).IsResetable());
  }

  void Reset() override {  //
    ITERTOOL_DETAIL_FoldExpr_ForEach(std::get<Index>(data_).Reset());
  }

 private:
  template <typename U, size_t Idx>
  U MaybeGetNow() {
    auto it = std::get<Idx>(data_);
    return it.MaybeGetNow(std::get<Idx>(default_value_));
  }

  IteratorTuple<T...> data_;
  std::tuple<T...> default_value_;
};

template <typename... T, size_t... Index>
class ZipTupleIterable<TypeSequence<T...>, std::index_sequence<Index...>>
    : public Iterable<ZipTupleItem<T...>> {
 public:
  using ItemType = ZipTupleItem<T...>;

  explicit ZipTupleIterable(IteratorTuple<T...> data) : data_(data) {}

  bool HasNext() override {  //
    return ITERTOOL_DETAIL_FoldExpr_AllOf(std::get<Index>(data_).HasNext());
  }

  ItemType GetNow() override {  //
    return ItemType(std::get<Index>(data_).GetNow()...);
  }

  void Advance() override {  //
    ITERTOOL_DETAIL_FoldExpr_ForEach(std::get<Index>(data_).Advance());
  }

  bool IsFinite() override {  //
    return ITERTOOL_DETAIL_FoldExpr_AnyOf(std::get<Index>(data_).IsFinite());
  }

  bool IsResetable() override {  //
    return ITERTOOL_DETAIL_FoldExpr_AllOf(std::get<Index>(data_).IsResetable());
  }

  void Reset() override {  //
    ITERTOOL_DETAIL_FoldExpr_ForEach(std::get<Index>(data_).Reset());
  }

 private:
  IteratorTuple<T...> data_;
};

template <template <class, class> class MaybeZipLongestIter, typename... T>
using MakeZipTupleIterable =
    MaybeZipLongestIter<TypeSequence<T...>, std::index_sequence_for<T...>>;

}  // namespace detail

template <typename... T>
Iterator<ZipTupleItem<T...>> Zip(Iterator<T>... args) {
  using IterType = detail::MakeZipTupleIterable<detail::ZipTupleIterable, T...>;
  return MakeIterator<IterType>(detail::IteratorTuple<T...>(args...));
}

template <typename... T>
Iterator<ZipTupleItem<T...>> ZipLongest(
    Iterator<T>... args,
    std::tuple<T...> default_value = std::make_tuple(T()...)) {
  using IterType =
      detail::MakeZipTupleIterable<detail::ZipLongestTupleIterable, T...>;
  return MakeIterator<IterType>(detail::IteratorTuple<T...>(args...),
                                default_value);
}

}  // namespace itertool
